"""
https://zhuanlan.zhihu.com/p/363153781
语音识别中的MFCC 与 Fbank特征的基本原理和python实现
https://www.cnblogs.com/LXP-Never/p/10918590.html
语音信号的梅尔频率倒谱系数(MFCC)的原理讲解及python实现
"""
import numpy
import scipy.io.wavfile
from scipy.fftpack import dct
import os
import sys
from python_ai.common.xcommon import *
import seaborn as sbn
import scipy

BASE_DIR, FILE_NAME = os.path.split(__file__)

# 一、预处理
# voice_path = r"C:\BAWEIJIAOYU\temp\src\CV\1904A_C5\语音识别(stop_left_right)\1904A_voice_data\stop/zsn-stop1.wav"
path = '../../../../../large_data/audio/zsn-stop1.wav'
voice_path = os.path.join(BASE_DIR, path)

sample_rate, signal = scipy.io.wavfile.read(voice_path)
# signal = signal[0:int(3.5 * sample_rate)]  # 我们只取前3.5s

print("signal:", np.shape(signal))  # (31744, 2)
signal = signal[:, 0]  # ATTENTION 左声道
print("signal:", np.shape(signal))  # (31744, )
print("sample_rate:", sample_rate)  # 16000

import matplotlib.pyplot as plt

spr = 2
spc = 1
spn = 0
plt.figure(figsize=[6, 6])

spn += 1
plt.subplot(spr, spc, spn)
plt.plot(signal)
plt.title("orignal_signal")
plt.grid()

# （1）、预加重（Pre-Emphasis）
pre_emphasis = 0.97
print('signal[0]', np.shape(signal[0]))  # (2, )
print('signal[1:]', np.shape(signal[1:]))  # (31743, )
print('signal[:-1]', np.shape(signal[:-1]))  # (31743, )
print('pre_emphasis * signal[:-1]', np.shape(pre_emphasis * signal[:-1]))  # (31743, )
emphasized_signal = numpy.append(signal[0], signal[1:] - pre_emphasis * signal[:-1])  # ATTENTION If `axis` is not given, both `arr` and `values` are flattened before use.
print('emphasized_signal', np.shape(emphasized_signal))  # (31744, )
print_numpy_ndarray_info(emphasized_signal, 'emphasized_signal')

spn += 1
plt.subplot(spr, spc, spn)
plt.plot(emphasized_signal)
plt.title("emphasized_signal")
plt.grid()

plt.show()

#（2）、分帧（Framing）
frame_size = 0.025  # 1/40, 25ms
frame_stride = 0.01  # 1/100, 10ms, 帧移
frame_length, frame_step = frame_size * sample_rate, frame_stride * sample_rate  # 从秒转换为采样点
signal_length = len(emphasized_signal)
frame_length = int(round(frame_length))
frame_step = int(round(frame_step))
print('signal_length', signal_length)  # 31744
print('frame_length', frame_length)  # 400
print('frame_step', frame_step)  # 160
# 确保我们至少有1帧
num_frames = int(numpy.ceil(float(numpy.abs(signal_length - frame_length)) / frame_step))
print('num_frames', num_frames)  # 196

pad_signal_length = num_frames * frame_step + frame_length
print('pad_signal_length', pad_signal_length)  # 31760
z = numpy.zeros((pad_signal_length - signal_length))
# 填充信号，确保所有帧的采样数相等，而不从原始信号中截断任何采样
pad_signal = numpy.append(emphasized_signal, z)
print_numpy_ndarray_info(pad_signal, 'pad_signal')  # (31760,)

idx01 = numpy.tile(
    numpy.arange(0, frame_length),
    (num_frames, 1)
)
print(idx01)  # (169, 400)
print_numpy_ndarray_info(idx01, 'idx01')
idx02 = numpy.tile(
    numpy.arange(0, num_frames * frame_step, frame_step),
    (frame_length, 1)
).T
print('W', int(np.floor(num_frames * frame_step / frame_step)))  # 0到169*160-1，共169*160个数，除以160，正好是169
print('H', frame_length)
print(idx02)
print_numpy_ndarray_info(idx02, 'idx02')

indices = numpy.tile(  # ATTENTION tile
    numpy.arange(0, frame_length),
    (num_frames, 1)
) + numpy.tile(
    numpy.arange(0, num_frames * frame_step, frame_step),
    (frame_length, 1)
).T
print(indices)
print_numpy_ndarray_info(indices, 'indices')  # (196, 400)

frames = pad_signal[indices.astype(numpy.int32, copy=False)]
print_numpy_ndarray_info(frames, 'frames')  # (196, 400)

# （3）、加窗（Window）
ham = numpy.hamming(frame_length)
print_numpy_ndarray_info(ham, 'hamming')  # (400, )
plt.plot(ham)
plt.title('hamming window')
plt.show()

# The Broadcasting Rule
# In order to broadcast,
# the size of the trailing axes for both arrays in an operation must either be the same size or one of them must be one.
frames *= ham  # ATTENTION 广播 (196, 400) * (400, ) => (196, 400)

# frames *= 0.54 - 0.46 * numpy.cos((2 * numpy.pi * n) / (frame_length - 1))  # 内部实现
print_numpy_ndarray_info(frames, 'frames after window')  # (196, 400)

# 二、FFT（Fourier-Transform）

print(numpy.fft.rfft, numpy.fft.fft2)
# np.fft.rfft Compute the one-dimensional discrete Fourier Transform for real input.
# np.fft.fft2 Compute the 2-dimensional discrete Fourier Transform
# If `n` is even, the length of the transformed axis is ``(n/2)+1``.
#         If `n` is odd, the length is ``(n+1)/2``.
NFFT = 512  # 512=>512/2+1=257 时域512=>频域257 FFT要求时域2^n
print_numpy_ndarray_info(frames, 'frames')
mag_frames = numpy.absolute(numpy.fft.rfft(frames, NFFT))   # fft的幅度(magnitude)
print_numpy_ndarray_info(mag_frames, 'mag_frames')  # (196, 257)

spr = 1
spc = 2
spn = 0
plt.figure(figsize=[12, 6])

spn += 1
plt.subplot(spr, spc, spn)
sbn.heatmap(mag_frames)
plt.title('Result of np.fft.rfft')

# 三、功率谱（Power Spectrum）
print_numpy_ndarray_info(mag_frames ** 2, 'mag_frames ** 2')  # (196, 257)
pow_frames = ((1.0 / NFFT) * (mag_frames ** 2))  # 功率谱  # ATTENTION 广播 (-1)*(196, 257) = (196, 257)
print_numpy_ndarray_info(pow_frames, 'pow_frames')  # (196, 257)

spn += 1
plt.subplot(spr, spc, spn)
sbn.heatmap(pow_frames)
plt.title('Power Spectrum')
plt.show()

import numpy as np

# 四、滤波器组（Filter Banks）
nfilt = 40
low_freq_mel = 0
high_freq_mel = (2595 * np.log10(1 + (sample_rate / 2) / 700))  # 将Hz转换为Mel  # ATTENTION sample_rate / 2 采样定理
print('high_freq_mel', high_freq_mel)
# 我们要做40个滤波器组，为此需要42个点，这意味着在们需要low_freq_mel和high_freq_mel之间线性间隔40个点
mel_points = np.linspace(low_freq_mel, high_freq_mel, nfilt + 2)  # 使得Mel scale间距相等
hz_points = (700 * (10 ** (mel_points / 2595) - 1))  # 将Mel转换回-Hz
# bin = sample_rate/NFFT    # frequency bin的计算公式
# bins = hz_points/bin=hz_points*NFFT/ sample_rate    # 得出每个hz_point中有多少frequency bin
bins = np.floor((NFFT + 1) * hz_points / sample_rate)
print_numpy_ndarray_info(bins, 'bins')
y1y2_bins = np.tile([[100], [-100]], len(bins))
print_numpy_ndarray_info(y1y2_bins, 'y1y2_bins')
plt.plot([bins, bins], y1y2_bins)
plt.title('bins')
plt.show()

fbank = np.zeros((nfilt, int(np.floor(NFFT / 2 + 1))))  # ATTENTION np.fft.rfft return NFFT => NFFT / 2 + 1
print_numpy_ndarray_info(fbank, 'fbank')  # (40, 257)

for m in range(1, nfilt + 1):
    f_m_minus = int(bins[m - 1])  # 左
    f_m = int(bins[m])  # 中
    f_m_plus = int(bins[m + 1])  # 右

    for k in range(f_m_minus, f_m):
        fbank[m - 1, k] = (k - bins[m - 1]) / (bins[m] - bins[m - 1])
    for k in range(f_m, f_m_plus):
        fbank[m - 1, k] = (bins[m + 1] - k) / (bins[m + 1] - bins[m])
filter_banks = np.dot(pow_frames, fbank.T)  # (196, 257) dot (257, 40) => (196, 40)
filter_banks = np.where(filter_banks == 0, np.finfo(float).eps, filter_banks)  # 数值稳定性
filter_banks = 20 * np.log10(filter_banks)  # dB

print('filter_banks', filter_banks.shape)

spr = 2
spc = 2
spn = 0
plt.figure(figsize=[8, 8])

spn += 1
plt.subplot(spr, spc, spn)
sbn.heatmap(filter_banks.T)
plt.title("filter_banks")

# 五、梅尔频率倒谱系数（MFCCs）
# 离散余弦变换（DCT）
num_ceps = 13
mfcc = scipy.fftpack.dct(filter_banks, type=2, axis=1, norm='ortho')[:, 1 : (num_ceps + 1)] # 保持在2-13

cep_lifter = 13
(nframes, ncoeff) = mfcc.shape
n = numpy.arange(ncoeff)
lift = 1 + (cep_lifter / 2) * numpy.sin(numpy.pi * n / cep_lifter)
mfcc *= lift

print("mfcc: ", mfcc.shape)

spn += 1
plt.subplot(spr, spc, spn)
sbn.heatmap(mfcc.T)
plt.title("mfcc")

# 六、均值归一化（Mean Normalization）
filter_banks -= (numpy.mean(filter_banks, axis=0) + 1e-8)
spn += 1
plt.subplot(spr, spc, spn)
sbn.heatmap(filter_banks.T)
plt.title("norm_filter_banks")

mfcc -= (numpy.mean(mfcc, axis=0) + 1e-8)
spn += 1
plt.subplot(spr, spc, spn)
sbn.heatmap(mfcc.T)
plt.title("norm_mfcc")
plt.show()
